import click
from collections import defaultdict
import math
import re
import decimal

genomesize = {} 
def dealsize():
    taxSizeList = '/home/yong_sun/bin/workflow/tax2size.list'
    with open(taxSizeList) as t2s:
        d = {}
        for i in t2s:
            tax, size = i.split()
            d[tax] = int(size)
            genomesize[tax] = size
        return d

def dealname():
    en2cn = '/home/yong_sun/bin/workflow/en2cn.list'
    with open(en2cn) as ec:
        d = {}
        for i in ec:
            en, cn = i.split('\t')
            d[en] = cn.strip()
        return d

@click.command()
@click.option('-nt', '--ntfile', required=True, type=str, help="input nt result file")
@click.option('-info', '--infofile', required=True, type=str, help="input kraken genus result file")
@click.option('-o', '--out', required=True, type=str, help="out nt result file")
@click.option('-co','--cofile',required=True, type=str, help="coverage file")
def main(ntfile, infofile, out, cofile):
    taxSize = dealsize()
    name = dealname()
    d_coverage=defaultdict(lambda:"0")
    ntCount = {}
    scalentCount = {}
    scaleSum = defaultdict(int)
    ntCountU = {}
    ntCountG = {}
    Info = []
    Ge = []
    Genus = defaultdict(int)
    sampleid = ntfile.strip().split('.')[0].split('-')[0]

    with open(cofile,encoding='utf-8') as fcoverage:
        for i in fcoverage:
            c_name,c_taxid,c_coverage,c_size=i.strip().split('\t')
            d_coverage[str(c_taxid)] =str(c_coverage)

    with open(infofile, encoding='utf-8') as f:
        for i in f:
            if i.strip().split('\t')[2] == '脑脊液' or i.strip().split('\t')[2] == '血液':
                Info.append(i.split('\t')[0])

    with open(ntfile) as ntf, open(out, 'w', encoding='gbk') as outf:
        ntf.readline()

        if sampleid in Info:
            filters = ['sp.', 'Rhizobium', 'phage', 'uncultured', 'Malassezia restricta', 'Sphingomonas', 'Comamonas', 'Micrococcus', \
                                'Mycobacterium leprae', 'Agrobacterium', 'Microbacterium', 'Escherichia virus', 'Burkholderia vietnamiensis', \
                                'Pandoraea pnomenusa', 'Acinetobacter guillouiae', 'Microbacterium laevaniformans', 'Yarrowia lipolytica', \
                                'Herbaspirillum','Xanthomonas','Methylobacterium','Cupriavidus','Delftia','Spirometra erinaceieuropaei']
        else:
            filters = ['sp.', 'Rhizobium', 'phage', 'uncultured', 'Malassezia restricta', 'Sphingomonas', 'Comamonas', 'Micrococcus', \
                            'Mycobacterium leprae', 'Agrobacterium', 'Microbacterium', 'Escherichia virus', 'Burkholderia vietnamiensis', \
                            'Pandoraea pnomenusa', 'Acinetobacter guillouiae', 'Microbacterium laevaniformans', 'Yarrowia lipolytica','Spirometra erinaceieuropaei']

        for i in ntf:
            info = i.split('\t')
            if info[0].find('undef') == -1 and int(info[9]) != 0 and all([i not in info[7] for i in filters]):
                ntTaxid = info[7].split('|')[0]
                if ntTaxid in taxSize and info[7].find('Ralstonia') == -1 and info[7].find('Erwinia amylovora') == -1:
                    ntCount[f'{info[0]}@{info[2]}@{info[7]}'] = info[9]
                    scalentCount[f'{info[0]}@{info[2]}@{info[7]}'] = int(info[9]) / taxSize[ntTaxid]
#                    if info[9]==0:continue 
                    scaleSum[info[0]] += int(info[9]) / taxSize[ntTaxid]
                else:
                    ntCountU[f'{info[0]}@{info[2]}@{info[7]}'] = info[9]
                ntCountG[f'{info[0]}@{info[2]}@{info[7]}'] = info[4] if int(info[4]) != 0 else info[3]
                if info[2] not in Ge:
                    Ge.append(info[2])
                    Genus[info[0]] += int(info[4])
            if info[7].find('Human betaherpesvirus 6A') > -1 or info[7].find('Human betaherpesvirus 6B') > -1:
                ntTaxid = info[7].split('|')[0]
                if ntTaxid in taxSize:
                    ntCount[f'{info[0]}@{info[2]}@{info[7]}'] = info[8]
                    scalentCount[f'{info[0]}@{info[2]}@{info[7]}'] = int(info[8]) / taxSize[ntTaxid]
                    scaleSum[info[0]] += int(info[8]) / taxSize[ntTaxid]
                    ntCountG[f'{info[0]}@{info[2]}@{info[7]}'] = info[4] if int(info[4]) != 0 else info[3]
                    if info[2] not in Ge:
                        Ge.append(info[2])
                        Genus[info[0]] += int(info[4])
            if info[7].find('Aspergillus flavus') > -1:
                ntTaxid = info[7].split('|')[0]
                flag = 0
                for j in info[11].split(';'):
                    if j.find('Aspergillus') == -1:
                        flag = 1
                if flag == 0:
                    flavus_num = int(info[8])
                    for j in info[11].split(';'):
                        if j.find('Aspergillus oryzae') == -1 or j.find('Aspergillus parasiticus') == -1:
                            a,n = j.split('=')
                            flavus_num -= int(n)
                else:
                    flavus_num = int(info[9])
                    for j in info[11].split(';'):
                        if j.find('Aspergillus oryzae') > -1 or j.find('Aspergillus parasiticus') > -1:
                            a,n = j.split('=')
                            flavus_num += int(n)
                if flavus_num > int(info[8]) or flavus_num < 0:
                    if int(info[8]) > int(info[4]):
                        flavus_num = int(info[4])
                    else:
                        flavus_num = int(info[8])
                ntCount[f'{info[0]}@{info[2]}@{info[7]}'] = str(flavus_num)
                scalentCount[f'{info[0]}@{info[2]}@{info[7]}'] = int(flavus_num) / taxSize[ntTaxid]
                scaleSum[info[0]] += int(flavus_num) / taxSize[ntTaxid]
                ntCountG[f'{info[0]}@{info[2]}@{info[7]}'] = info[4] if int(info[4]) != 0 else info[3]
                if info[2] not in Ge:
                    Ge.append(info[2])
                    Genus[info[0]] += int(info[4])
            if info[7].find('Streptococcus pneumoniae') > -1:
                ntTaxid = info[7].split('|')[0]
                pneumoniae_num = int(info[9])
                for j in info[11].split(';'):
                    if j.find('Streptococcus mitis') > -1:
                        a,n = j.split('=')
                        pneumoniae_num += int(n)
                ntCount[f'{info[0]}@{info[2]}@{info[7]}'] = str(pneumoniae_num)
                scalentCount[f'{info[0]}@{info[2]}@{info[7]}'] = int(pneumoniae_num) / taxSize[ntTaxid]
                scaleSum[info[0]] += int(pneumoniae_num) / taxSize[ntTaxid]
                ntCountG[f'{info[0]}@{info[2]}@{info[7]}'] = info[4] if int(info[4]) != 0 else info[3]
                if info[2] not in Ge:
                    Ge.append(info[2])
                    Genus[info[0]] += int(info[4])
        outf.write(f'D\tG\tS\tCN_name\tsReads\tgReads\tRelativeAbundance\tGenustRelativeAbundance\tgenomesize\tcoverage(%)\n')
        for k, v in ntCount.items():
            info = k.split('@')
            tax = info[2].split('|')[0]
            ratio = int(v) / int(ntCountG[k])
            coverage=str(d_coverage[str(tax)])#added
             
#            print(type(decimal.Decimal(scalentCount[k])))
#            print(scalentCount[k])
            GscaleRB = f'{int(ntCountG[k]) / Genus[info[0]] * 100}'
            scaleRB = f'{scalentCount[k] / scaleSum[info[0]] * 100 * ratio}'
#            scaleRB = f'{decimal.Decimal(scalentCount[k]) / decimal.Decimal(str(scaleSum[info[0]]))}'
            flag = 0
            if info[2].split('|')[-1] in name :
                flag = info[2].split('|')[-1]
            if flag:
                outf.write('\t'.join((*info, name[flag], v, ntCountG[k], str(scaleRB), str(GscaleRB), genomesize[tax],coverage))+'\n')
            else:
                outf.write('\t'.join((*info, '-', v, ntCountG[k], str(scaleRB), str(GscaleRB), genomesize[tax],coverage))+'\n')
        for k, v in ntCountU.items():
            info = k.split('@')
            tax2=info[2].split('|')[0]
            coverage=str(d_coverage[str(tax2)])#added
            flag = 0
            if info[2].split('|')[-1] in name :
                flag = info[2].split('|')[-1]
            if flag:
                if info[2] == '552|Erwinia amylovora':
                    outf.write('\t'.join(('-', '-', '内标', '内标', v, ntCountG[k], '-', '-', genomesize[tax],coverage))+'\n')
                else:
                    outf.write('\t'.join((*info, name[flag], v, ntCountG[k], '-', '-', genomesize[tax],coverage))+'\n')
            else:
                if info[2] == '552|Erwinia amylovora':
                    outf.write('\t'.join(('-', '-', '内标', '内标', v, ntCountG[k], '-', '-', genomesize[tax],coverage))+'\n')
                else:
                    outf.write('\t'.join((*info, '-', v, ntCountG[k], '-', '-', genomesize[tax],coverage))+'\n')

if __name__ == "__main__":
    main()
